package org.github.manx98.jlibhac.tools.fssystem;

import org.github.manx98.jlibhac.common.Utilities;
import org.github.manx98.jlibhac.common.math.UnsignedNumber;
import org.github.manx98.jlibhac.common.math.Uint64;
import org.github.manx98.jlibhac.util.Alignment;

import javax.crypto.Cipher;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.ShortBufferException;
import javax.crypto.spec.SecretKeySpec;
import java.io.IOException;
import java.nio.ByteOrder;
import java.security.InvalidKeyException;
import java.security.NoSuchAlgorithmException;

/**
 * @author Manx98
 * @version 1.0
 * @since 2022/10/3 15:40
 */
public class Aes128XtsTransform {
    private final Cipher key1;
    private final Cipher key2;
    private final boolean decrypting;
    private final int BlockSize = 128;
    private final int BlockSizeBytes = BlockSize / 8;
    private final byte[] cc = new byte[16];
    private final byte[] pp = new byte[16];
    private final byte[] t = new byte[16];

    public Aes128XtsTransform(byte[] key1, byte[] key2, boolean decrypting) {
        try {
            this.decrypting = decrypting;
            if (decrypting) {
                this.key1 = createCipher(key1, Cipher.DECRYPT_MODE);
                this.key2 = createCipher(key2, Cipher.ENCRYPT_MODE);
            } else {
                this.key1 = createCipher(key1, Cipher.ENCRYPT_MODE);
                this.key2 = createCipher(key2, Cipher.ENCRYPT_MODE);
            }
        } catch (NoSuchAlgorithmException | NoSuchPaddingException | InvalidKeyException e) {
            throw new RuntimeException(e);
        }
    }

    private Cipher createCipher(byte[] key, int mod) throws NoSuchPaddingException, NoSuchAlgorithmException, InvalidKeyException {
        Cipher instance = Cipher.getInstance("AES/ECB/NoPadding");
        instance.init(mod, new SecretKeySpec(key, "AES"));
        return instance;
    }

    public void transformBlock(byte[] buffer, int offset, int count, long sector) throws IOException {
        int lim;

        /* get number of blocks */
        int m = count >> 4;
        int mo = count & 15;
        int alignedCount = Alignment.alignUp(count, BlockSizeBytes);

        /* for i = 0 to m-2 do */
        if (mo == 0) {
            lim = m;
        } else {
            lim = m - 1;
        }
        byte[] tweak = new byte[alignedCount];
        fillArrayFromSector(tweak, sector);
        try {
            key2.update(tweak, 0, 16, tweak, 0);
            fillTweakBuffer(tweak);
            if (lim > 0) {
                Utilities.xorArrays(buffer, offset, lim * 16, tweak);
                key1.update(buffer, offset, lim * 16, buffer, offset);
                Utilities.xorArrays(buffer, offset, lim * 16, tweak);
            }
            if (mo > 0) {
                if (decrypting) {
                    System.arraycopy(tweak, lim * 16 + 16, cc, 0, 16);
                    doTweakCrypt(buffer,offset,mo, cc, pp);
                } else {
                    doTweakCrypt(buffer,offset, mo, pp, cc);
                }
            }
        } catch (ShortBufferException e) {
            throw new RuntimeException(e);
        }
    }

    public void doTweakCrypt(byte[] buffer, int offset, int mo, byte[] pp, byte[] cc) throws ShortBufferException {
        /* CC = tweak encrypt block m-1 */
        tweakCrypt(buffer, offset, cc, 0, t);

        /* Cm = first ptlen % 16 bytes of CC */
        int i;
        for (i = 0; i < mo; i++) {
            pp[i] = buffer[16 + i + offset];
            buffer[16 + i + offset] = cc[i];
        }

        for (; i < 16; i++) {
            pp[i] = cc[i];
        }

        /* Cm-1 = Tweak encrypt PP */
        tweakCrypt(pp, 0, buffer, offset, t);
    }

    /**
     * Performs the Xts TweakCrypt operation
     */
    private void tweakCrypt(byte[] inputBuffer, int inputOffset, byte[] outputBuffer, int outputOffset, byte[] t) throws ShortBufferException {
        for (int x = 0; x < 16; x++) {
            outputBuffer[x + outputOffset] = (byte) (inputBuffer[x + inputOffset] ^ t[x]);
        }

        key1.update(outputBuffer, outputOffset, 16, outputBuffer, outputOffset);

        for (int x = 0; x < 16; x++) {
            outputBuffer[x + outputOffset] = (byte) (outputBuffer[x + outputOffset] ^ t[x]);
        }
        multiplyByX(t);
    }

    /**
     * Multiply by x
     */
    private void multiplyByX(byte[] i) {
        byte t = 0, tt = 0;

        for (int x = 0; x < 16; x++) {
            tt = (byte) (i[x] >> 7);
            i[x] = (byte) (((i[x] << 1) | t) & 0xFF);
            t = tt;
        }

        if (tt > 0) {
            i[0] ^= 0x87;
        }
    }

    private void fillArrayFromSector(byte[] value, long sector) {
        for (int i = 0xF; i >= 0; i--) {
            value[i] = (byte) sector;
            sector >>= 8;
        }
    }

    public static void fillTweakBuffer(byte[] buffer) throws IOException {
        UnsignedNumber b = new Uint64(buffer, 0);
        UnsignedNumber a = new Uint64(buffer, 8);
        UnsignedNumber magic = new Uint64(0x87);
        for (int i = 2; i < buffer.length / 8; i += 2) {
            UnsignedNumber tt = new Uint64(a.longValue() >> 63).and(magic);
            a = a.shiftLeft(1).or(b.shiftRight(63));
            b = b.shiftLeft(1).xor(tt);
            a.toByteArray(buffer, (i + 1) * 8, ByteOrder.LITTLE_ENDIAN);
            b.toByteArray(buffer, i * 8, ByteOrder.LITTLE_ENDIAN);
        }
    }
}
